{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://mp.weixin.qq.com/s?__biz=MzkyMDMxNDkxOA==&mid=2247484455&idx=1&sn=0f0994f5e25c344572190380315d6b87&chksm=c195f16ef6e278783c1418c0524d01b15aa66b3344655453c45e82f72d562f816bf7abc022c1&scene=178&cur_album_id=3447886516519747585#rd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 振动升降指标ASI策略.py\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from datetime import datetime\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_hist_k_data(code,start_date,end_date,frequency='d')->pd.DataFrame:\n",
    "    \"\"\"\n",
    "    获取历史K线数据\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    import baostock as bs\n",
    "    bs.login()\n",
    "    rs = bs.query_history_k_data_plus(code,\"date,code,open,high,low,close,preclose,volume,amount,pctChg\",start_date,end_date,frequency=frequency)\n",
    "    data_list = []\n",
    "    while (rs.error_code == '0') & rs.next():\n",
    "        # 获取一条记录，将记录合并在一起\n",
    "        data_list.append(rs.get_row_data())\n",
    "    result = pd.DataFrame(data_list, columns=rs.fields)\n",
    "    result.open = result.open.astype(float)\n",
    "    result.high = result.high.astype(float)\n",
    "    result.low = result.low.astype(float)\n",
    "    result.close = result.close.astype(float)\n",
    "    result.date= pd.to_datetime(result.date)\n",
    "    result.set_index('date',inplace=True)\n",
    "    bs.logout()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获取股票数据\n",
    "stock_symbol = 'sz.000001'\n",
    "start_date = '2005-01-01'\n",
    "end_date = '2023-01-01'\n",
    "data = get_hist_k_data(stock_symbol, start_date, end_date)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 计算DMI指标\n",
    "def calculate_dmi(data, period=14):\n",
    "    high = data['high']\n",
    "    low = data['low']\n",
    "    close = data['close']\n",
    "    tr = pd.DataFrame({\n",
    "        'tr': (high - low).abs(),\n",
    "        'tr_high': (high - close.shift(1)).abs(),\n",
    "        'tr_low': (close.shift(1) - low).abs()\n",
    "    }).max(axis=1)\n",
    "    tr = tr.rolling(window=period).sum()\n",
    "    dm_plus = (high - high.shift(1)).clip(lower=0).rolling(window=period).sum()\n",
    "    dm_minus = (low.shift(1) - low).clip(lower=0).rolling(window=period).sum()\n",
    "    di_plus = 100 * dm_plus / tr\n",
    "    di_minus = 100 * dm_minus / tr\n",
    "    dx = (abs(dm_plus - dm_minus) / (dm_plus + dm_minus)) * 100\n",
    "    adx = dx.rolling(window=period).mean()\n",
    "    return pd.DataFrame({\n",
    "        'DI+': di_plus,\n",
    "        'DI-': di_minus,\n",
    "        'ADX': adx\n",
    "    })\n",
    "\n",
    "# 设置ADX阈值\n",
    "adx_threshold = 25\n",
    "\n",
    "# 生成交易信号\n",
    "df = calculate_dmi(data)\n",
    "df['Signal'] = 0\n",
    "df['Position'] = 0\n",
    "\n",
    "# 当+DI交叉-DI时买入\n",
    "df['Signal'][df['DI+'] > df['DI-']] = 1\n",
    "# 当-DI交叉+DI时卖出\n",
    "df['Signal'][df['DI-'] > df['DI+']] = -1\n",
    "\n",
    "# 计算持仓\n",
    "df['Position'] = df['Signal'].diff()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data.join(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 回测策略\n",
    "def backtest_strategy(data):\n",
    "    data['Strategy_Returns'] = data['Position'].shift(1) * data['close'].pct_change()\n",
    "    data['Cumulative_Returns'] = (1 + data['Strategy_Returns']).cumprod()\n",
    "    return data\n",
    "\n",
    "data = backtest_strategy(data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 绘制策略的累计收益\n",
    "plt.figure(figsize=(14, 7))\n",
    "plt.plot(data['Cumulative_Returns'], label='AO Strategy')\n",
    "plt.title('AO Strategy Cumulative Returns')\n",
    "plt.xlabel('Date')\n",
    "plt.ylabel('Cumulative Returns')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
